#!/usr/bin/python3
import pickle
import cv2
import numpy as np
import sys

import Camera

def test_calib():
    print("test calib")

    cam_idx_list, calib_mtxs, calib_dists = Camera.load_camera_params()
    
    cp = Camera.get_capture()
    idx = 0
    while True:
        i = cam_idx_list[idx]
        idx += 1
        idx %= 4
        print(i)
        cp.open(i)
        cp.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M','J','P','G'))
        cp.set(cv2.CAP_PROP_FRAME_WIDTH, Camera.WIDTH)
        cp.set(cv2.CAP_PROP_FRAME_HEIGHT, Camera.HEIGHT)

        if not cp.isOpened():
            continue
        while 1:
            _, frame = cp.read()
            if not _:
                continue
            cv2.imshow("cam" + str(frame.shape), frame)
            #cv2.imshow("calib", Camera.camera_calibration(frame, calib_mtxs[cam_idx_list[idx], calib_dists[cam_idx_list[idx]))
            key = cv2.waitKey(1)
            if key == ord('q'):
                exit()
            if key == ord('c'):
                break

def calibrate_camera(images):    
    criteria = (cv2.TERM_CRITERIA_MAX_ITER | cv2.TERM_CRITERIA_EPS, 30, 0.001)

    objp = np.zeros((6 * 8, 3), np.float32)
    objp[:, :2] = np.mgrid[0:6, 0:8].T.reshape(-1, 2)

    obj_points = []
    img_points = []

    for i in images:
        img = i
        # cv2.imshow('img',img)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        size = gray.shape[::-1]
        ret, corners = cv2.findChessboardCorners(gray, (6, 8), None)
        print(ret)

        if ret:
            obj_points.append(objp)
            corners2 = cv2.cornerSubPix(gray, corners, (5, 5), (-1, -1), criteria)
            
            if [corners2]:
                img_points.append(corners2)
            else:
                img_points.append(corners)

            cv2.drawChessboardCorners(img, (6, 8), corners, ret)            
            cv2.imshow('img', img)
            cv2.waitKey(1000)

    print(len(img_points))
    cv2.destroyAllWindows()

    # calib
    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(obj_points, img_points, size, None, None)

    print("mtx:\n", mtx)
    print("dist:\n", dist)
    img = images[0]
    h, w = img.shape[:2]
    newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx,dist,(w,h),1,(w,h))
    dst = cv2.undistort(img,mtx,dist,None,newcameramtx)
    x,y,w,h = roi
    dst1 = dst[y:y+h,x:x+w]
    # Calc Error
    tot_error = 0
    for i in range(len(obj_points)):
        imgpoints2, _ = cv2.projectPoints(obj_points[i], rvecs[i], tvecs[i], mtx, dist)
        error = cv2.norm(img_points[i],imgpoints2, cv2.NORM_L2)/len(imgpoints2)
        tot_error += error

    print ("total error: ", tot_error/len(obj_points))
    cv2.imshow("test",dst1)
    cv2.waitKey(0)

    return mtx, dist


'''
# 生成含有标定板的图像的序列
imagesfilename = glob.glob("*.jpg")
images = []
for fname in imagesfilename:
    temp=cv2.imread(fname)
    images.append(temp)
'''

if __name__ == "__main__":
    cam_idx_list = []
    mtxs = {}
    dists = {}
    CP_OPEN = 0
    with open("../configs/cam_idx.bin", 'rb') as cam_idx_file:
        cam_idx_list = pickle.load(cam_idx_file)
    try:
        with open("../configs/cam_calib_mtx.bin", 'rb') as cam_calib_file:
            mtxs = pickle.load(cam_calib_file)
        with open("../configs/cam_calib_dist.bin", 'rb') as cam_calib_file:
            dists = pickle.load(cam_calib_file)
    except Exception as e:
        print(e)
        print("Problems occurred, start from scratch.")
        mtxs = {}
        dists = {}
    
    if len(sys.argv) > 1 and sys.argv[1] == 'test':
        test_calib()
        exit()

    cap = Camera.get_capture()
    cap.open(cam_idx_list[CP_OPEN])
    images = []

    while True:
        _, frame = cap.read()
        cv2.imshow("Calib" + str(frame.shape), frame)
        key = cv2.waitKey(1)
        if key == ord('q'):
            break
        if key == ord('c'):
            CP_OPEN += 1
            CP_OPEN %= 4
            cap.open(cam_idx_list[CP_OPEN])
            images.clear()
            continue
        if key == ord('s'):
            images.append(frame.copy())
            print("capture image ", len(images))
        if key == ord('o'):
            if len(images) >= 2:
                mtxs[CP_OPEN], dists[CP_OPEN] = calibrate_camera(images)
                with open("../configs/cam_calib_mtx.bin", 'wb') as cam_calib_file:
                    pickle.dump(mtxs, cam_calib_file)
                    print("saved to configs/cam_calib_mtx.bin")
                with open("../configs/cam_calib_dist.bin", 'wb') as cam_calib_file:
                    pickle.dump(dists, cam_calib_file)
                    print("saved to configs/cam_calib_dist.bin")
                images.clear()
            else:
                print("Too few image. More than 2 shots is recommened.")
